# -*- coding: utf-8 -*-
"""
This file is part of tptsig
Copyright © 2011 Telecom ParisTech, TSI
Auteur(s) : Liutkus Antoine
tptsig is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Lesser General Public License for more details.
You should have received a copy of the GNU LesserGeneral Public License
along with this program. If not, see <http://www.gnu.org/licenses/>."""


import numpy as np
from scipy import *
from hdf5buffer import  *
import wave


def wavinfo(filename):
	"""returns metadata properties of a file"""
	file = wave.open(filename, 'r');
	
	nChans = file.getnchannels();
	fs = file.getframerate();
	len = file.getnframes();
	sampleWidth = file.getsampwidth();
	file.close();
	return len, nChans, fs, sampleWidth;

def wavwrite(signal, fs, destinationFile, nbytes = 2):
	"""writes a wavfile into a file"""
	batchSize = 500000
	file = wave.open(destinationFile ,'wb')
	file.setparams((signal.data.shape[1], nbytes, fs, signal.data.shape[0], 'NONE', 'not compressed'));
	pos = 0
	len = signal.data.shape[0]
	
	while pos<len:
		batchSize = min(batchSize, len - pos)
		tempdata = minimum(abs(signal.data[pos:pos+batchSize,:].flatten()),0.98)*sign(signal.data[pos:pos+batchSize,:].flatten())
		dataVec = (2**(nbytes*8 - 1)*tempdata).astype(int16)
		values = dataVec.tostring()
		file.writeframes(values)
		pos +=batchSize
	file.close()
	print 'File written to ' + destinationFile


def isValid(filename):
	"""returns whether the file is valid."""
	if not filename:
		return False
	try:
		file = wave.open(filename, 'r');
		file.close()
		return True
	except  :
		return False

class Signal:
	"""Signal class. This class encapsulate basic operations with wave files.
	It can efficiently handle large wave files, perform Short Term Fourier Transform (STFT) 
	and compute its invers (iSTFT).
	the waveform is in the wave.data propery if available (wave is not None)
	the STFT is in the stft.data property if available (stft is not None)."""
	wave = None
	stft = None

	realAtom = None		#atom for real data (signals)
	complexAtom = None  #atom for complex data (STFT)
	
	fileName = None
	
	#The following property are set automatically when read is called
	shortFileName = None
	directoryName = None
	fs = None
	len = None
	nChans = None
	sampleWidth = None

	#STFT parameters
	nfft = None
	windowLen = 70
	overlap = 0.75
	
	
	def isValid(self):
		return isValid(self.fileName);
			
	def read(self):
		"""reads the filename file and set all the properties of the object"""
		if not self.isValid():
			print "invalid WAV file. Aborting"
		
		# filename properties
		(self.directoryName, self.shortFileName) = os.path.split(self.fileName)

		# metadata properties of a file
		self.len, self.nChans, self.fs, self.sampleWidth = wavinfo(self.fileName)
		self.nfft = round(self.fs/1000*70/2)*2;

		#erases previous data
		if self.wave is not None:
			del self.wave
		
		#creates buffer
		self.wave = hdf5Buffer(self.realAtom, (self.len,self.nChans))
		
		#reading data
		file = wave.open(self.fileName, 'r')
		pos = 0
		batchSize = 500000
		while pos<self.len:
			batchSize = min(batchSize, self.len - pos)
			str_bytestream = file.readframes(batchSize)
			tempData = fromstring(str_bytestream,'h')
			tempData = tempData.astype(float);
			tempData = tempData.reshape(batchSize, self.nChans);
			self.wave.data[pos:pos+batchSize,:] = tempData / 2**(8*self.sampleWidth - 1)
			pos +=batchSize
		file.close()
		
	def Help(self):
		print """ Signal class
		Handles wav files and common operations for signal processing : 
			* opening/reading/writing of multichannel audio wav
			* STFT, iSTFT with any window length and overlap ratio
		"""
	
	# Constructor
	def __init__(self, filename = None, nbytes=4, read = True, data = None, fs = None):
		if nbytes == 4 : 
			self.realAtom = 'float32'
			self.complexAtom = 'complex'
		elif nbytes == 8:
			self.realAtom = 'float64'
			self.complexAtom = 'complex'
		else:
			print "wrong number of bytes for data. Must be either 4 or 8.";

		self.fileName = filename

		if self.fileName and read:
			print 'reading WAV file'
			self.read()
		elif (data is not None) and (fs is not None):
			self.len, self.nChans = data.shape
			self.fs = fs
			self.sampleWidth = 16
			self.nfft = round(self.fs/1000*70/2)*2;
			self.wave = buffer((self.len,self.nChans));
			self.wave.data[:,:] = data[:,:]
			
			
	# ToString method
	def toString(self):
		print self.shortFileName , ' in ' , self.directoryName
		print "Channels :" , self.nChans
		print "Sample Frequency :" , self.fs , " Hz"        
		print "Duration :" , (self.len), " frames  or " , float(self.len)/float(self.fs) , ' seconds'
		print "Sample width :" , self.sampleWidth , ' bytes'
	
	
	def write(self, outputPath = '.', outputName = 'output.wav'):
		""" writes to an output file"""
		fileName=os.path.join(outputPath,outputName);
		wavwrite(self.wave, self.fs, fileName, self.sampleWidth)        ;

	
	def STFT(self):
		"""STFT (Short Time Fourier Transform) Computation"""
		self.nfft = int(round(self.nfft/2)*2)
		
		#building frames positions
		step = self.nfft-round(self.overlap*self.nfft);
		self.framesPositions = arange(0,self.wave.data.shape[0],step);
		self.framesPositions = self.framesPositions[nonzero(add(self.framesPositions,self.nfft-1) < self.wave.data.shape[0] )];
		self.framesPositions = append(self.framesPositions, self.framesPositions[-1] + step);
		self.framesPositions = array(self.framesPositions,dtype='int')
		self.nFrames = len(self.framesPositions);

		#building framed signal
		shape = (self.nfft,self.nFrames,self.nChans);
		win1D = hamming(self.nfft);
		win = outer(win1D, ones(self.nChans) );
		tempLength = self.framesPositions[-1]+self.nfft;     
		sWin = hdf5Buffer(self.realAtom, shape)

		#alocating memory for stft
		shape = (self.nfft/2 +1,self.nFrames,self.nChans);
		if self.stft is not None:
			del self.stft
		self.stft = hdf5Buffer(self.complexAtom, shape)

		batchSize = int(4000/self.nChans)
		batchSizeT = int(1000000/self.nChans)
		lastUpdate = 0
		firstSample = 0
		lastSample = min(self.len,batchSizeT)
		tempSourceSignal = self.wave.data[firstSample:lastSample,:]
		tempFramedSignal = zeros((self.nfft,batchSize,self.nChans))
		for index in range(0,self.nFrames-1):
			current_range = (arange(self.framesPositions[index],(self.framesPositions[index]+self.nfft))-firstSample).astype(int) 
			if not index%50:
				print 'STFT : frame ' + str(index) + '/' + str(self.nFrames)
			if index and (not index %batchSize):
				self.stft.data[:,lastUpdate:lastUpdate+batchSize,:] = np.fft.rfft(tempFramedSignal,axis = 0)            
				lastUpdate += batchSize
			tempFramedSignal[:,index%batchSize,:] = tempSourceSignal[current_range,:]*win;
			if self.framesPositions[index+1]+self.nfft > lastSample:
				firstSample = int(self.framesPositions[index+1])
				lastSample = int(min(self.len,firstSample+batchSizeT))
				tempSourceSignal = self.wave.data[firstSample:lastSample,:]
				
		self.stft.data[:,lastUpdate:self.stft.data.shape[1]-1,:] = np.fft.rfft(tempFramedSignal[:,:(index%batchSize+1),:],axis = 0)	
		lEnd = self.framesPositions[-1] + self.nfft  - self.len;
		finalZeroPadding = zeros((lEnd, self.nChans));   
		self.stft.data[:,self.stft.data.shape[1]-1,:] = np.fft.rfft(concatenate((self.wave.data[self.framesPositions[-1]:,:] , finalZeroPadding),axis=0)*win,axis=0);
		
	def iSTFT(self):
		"""computes inverse STFT"""
		def internalISTFT(S,framesPos):
			if not S.ndim:
				return;
			nfft = (S.shape[0]-1)*2
			win1D = hamming(self.nfft);
			squaredWin1D = win1D**2;
			
			shape = (nfft,S.shape[1],S.shape[2]);
			tempSignal = zeros(shape);
			batchSize = 2000;
			batchPos = range(0,S.shape[1],batchSize);
			for posBatch in batchPos[:-1]:
				tempSignal[:,posBatch:(posBatch+batchSize),:] = np.fft.irfft(S[:,posBatch:(posBatch+batchSize),:],axis = 0);
			tempSignal[:,batchPos[-1]:,:] = np.fft.irfft(S[:,batchPos[-1]:,:],axis = 0);            
			tempLength = framesPos[-1]-framesPos[0]+nfft;     
			sWeights = zeros((tempLength,S.shape[2]));
			s_temp=zeros((tempLength,S.shape[2]));
			for index in range(0,len(framesPos)):
				for index_chan in range(0, S.shape[2]):
					current_range = arange(framesPos[index]-framesPos[0],(framesPos[index]-framesPos[0]+nfft)).astype(int); 
					s_temp[current_range,index_chan] += tempSignal[:,index,index_chan]*win1D;
					sWeights[current_range,index_chan] += squaredWin1D;                
					
			s_temp /= sWeights;            
			return s_temp
		
		if self.stft is None:
			return;
		result = hdf5Buffer(self.realAtom,(self.len,self.nChans))
		batchSize = int(4000/self.nChans)
		batchPos = range(0,self.nFrames,batchSize);
		stopSample = -1
		print 'computing iSTFT'
		for posBatch in batchPos:
			lastFrame = min(self.nFrames,posBatch+batchSize)
			sigBatch = internalISTFT(self.stft.data[:,posBatch:lastFrame,:],self.framesPositions[posBatch:lastFrame])
			usefulLength = min(sigBatch.shape[0],self.len - self.framesPositions[posBatch])
			result.data[self.framesPositions[posBatch]:(self.framesPositions[posBatch]+usefulLength)] += sigBatch[0:usefulLength,:]
			if self.framesPositions[posBatch] < stopSample:
				result.data[self.framesPositions[posBatch]:stopSample]/=2
			stopSample = (self.framesPositions[posBatch]+usefulLength)
			
		return result;

